import numpy as np

# 创建迷宫地图
exit_coord = (3, 3)
row_n, col_n = 4, 4

maze = np.zeros((row_n, col_n)) - 1

# 走出迷宫奖励10个积分
maze[exit_coord] = 10

# 走到墙网格，扣除10个积分
maze[(0, 3)] = -10
maze[(1, 0)] = -10
maze[(1, 2)] = -10
maze[(2, 2)] = -10
maze[(3, 0)] = -10


# 定义动作集合
action_n = 4
actions = [0, 1, 2, 3]  # 上、下、左、右

# 定义参数
alpha = 0.1  # 学习率
gamma = 0.9  # 折扣因子
epsilon = 0.1  # ε-greedy策略的ε值


# 初始化Q表
Q = np.zeros((row_n, col_n, action_n))


# 进行Q-learning算法迭代更新
begin_cord = (0, 0)
max_reward_route = float("-inf")
for episode in range(200):
    # 初始化起始位置
    state = begin_cord
    route = [state]
    while state != exit_coord:  # 终止条件：到达终点位置
        tmp = actions.copy()
        # 排除一些可能
        if state[0] == 0:  # 不能向上
            tmp.remove(0)
        if state[1] == 0:  # 不能向左
            tmp.remove(2)
        if state[0] == row_n - 1:  # 不能向下
            tmp.remove(1)
        if state[1] == col_n - 1:  # 不能向右
            tmp.remove(3)

        # 选择动作
        if np.random.uniform() < epsilon:
            action = np.random.choice(tmp)  # ε-greedy策略，以一定概率随机选择动作
        else:
            action = np.argmax(Q[state[0], state[1], tmp])  # 选择Q值最大的动作
            action = tmp[action]

        # 执行动作，更新状态
        next_state = state
        if action == 0:  # 上
            next_state = (state[0] - 1, state[1])
        elif action == 1:  # 下
            next_state = (state[0] + 1, state[1])
        elif action == 2:  # 左
            next_state = (state[0], state[1] - 1)
        elif action == 3:  # 右
            next_state = (state[0], state[1] + 1)

        # 获取即时奖励
        reward = maze[next_state]

        # 更新Q值
        Q[state][action] = (1 - alpha) * Q[state][action] + alpha * (
            reward + gamma * np.max(Q[next_state])
        )

        # 更新状态
        state = next_state
        route.append(state)

    route_reward = sum(maze[state] for state in route)
    if max_reward_route < route_reward:
        max_reward_route = route_reward
        best_route = route.copy()
        print(f"episode: {episode}, 新发现最优路线：{best_route}")

    route.clear()
    cur_reward_route = 0
